{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "287fcd9a-bc4d-4b6f-b2d9-63675ecafc26",
   "metadata": {},
   "source": [
    "### 多线程版本 v2 两个并行方向\n",
    "1. 获取子目录后对每一个子目录并行获取子目录大小\n",
    "2. 获取子目录后对于每一个文件并行获取文件大小\n",
    "3. 使用scan_dir 手动控制递归和递归深度，如果需要更准确的结果需要更深的探索，性能相应的会损失\n",
    "\n",
    "### 目前看下来是最优的方案"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d4ccb344-305f-4b28-bec2-a683dfea0be0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "import concurrent.futures\n",
    "import threading\n",
    "import time\n",
    "\n",
    "import xlsxwriter\n",
    "\n",
    "LARGE_FILE_THRESHOLD = 200 * 1024 ** 2  # 200MB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "98ccb046-3521-4124-aaf1-a6bd8873ab26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取单个文件的大小\n",
    "def get_file_size(dirpath, files):\n",
    "    dir_size = 0        \n",
    "    with ThreadPoolExecutor(max_workers=8) as executor:   \n",
    "        futures = []\n",
    "        for file_name in files:\n",
    "            file_path = os.path.join(dirpath, file_name)\n",
    "            futures.append(executor.submit(os.path.getsize, file_path))\n",
    "            \n",
    "        for future in concurrent.futures.as_completed(futures):\n",
    "            try: \n",
    "                dir_size += future.result()\n",
    "            except FileNotFoundError:\n",
    "                pass\n",
    "    return dirpath, dir_size\n",
    "\n",
    "def scan_with_scandir(path, executor, futures, max_depth, current_depth=0):\n",
    "    if current_depth > max_depth:\n",
    "        return\n",
    "    try:\n",
    "        files = []\n",
    "        with os.scandir(path) as entries:\n",
    "            for entry in entries:\n",
    "                if entry.is_file():\n",
    "                    files.append(entry.name)\n",
    "                elif entry.is_dir(follow_symlinks=False):\n",
    "                    scan_with_scandir(entry.path, executor, futures, max_depth, current_depth + 1)\n",
    "        if files:\n",
    "            futures.append(executor.submit(get_file_size, path, files))\n",
    "    except (PermissionError, FileNotFoundError):\n",
    "        pass\n",
    "        \n",
    "# 获取文件夹下面所有文件的大小\n",
    "def get_dir_size(root,directory,max_workers=4,max_depth=8):\n",
    "    total_size = 0\n",
    "    base_stat={}\n",
    "    base_dir = os.path.join(root, directory)\n",
    "    \n",
    "    start_time = time.time()\n",
    "    with ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
    "        futures = []\n",
    "        scan_with_scandir(base_dir, executor, futures, max_depth)\n",
    "        \n",
    "        # 获取结果\n",
    "        for future in concurrent.futures.as_completed(futures):\n",
    "            result = future.result()\n",
    "            dir_size = result[1]\n",
    "            total_size += dir_size\n",
    "            if dir_size > LARGE_FILE_THRESHOLD:\n",
    "                base_stat[result[0]] = dir_size\n",
    "    end_time = time.time() \n",
    "    print(f\"分析文件夹 {base_dir} 用时{(end_time - start_time):.2f} 秒 ---------------------\")\n",
    "    return directory, total_size, base_stat\n",
    "\n",
    "\n",
    "# 扫描根文件夹, 拆分任务\n",
    "def scan_dirs(root):\n",
    "    print(f\"开始分析 Folder: {root} --------------------------------\")    \n",
    "    statMap = {}\n",
    "    total_size = 0\n",
    "    for entry in os.scandir(root):\n",
    "        try:\n",
    "            if entry.is_dir():  # 只保留子目录\n",
    "                size = get_dir_size(root, entry.name) \n",
    "                total_size += size[1]\n",
    "                statMap.update(size[2])\n",
    "                # print(root, entry.name)\n",
    "        except PermissionError:\n",
    "            continue\n",
    "         \n",
    "    return total_size,statMap\n",
    "\n",
    "# 生成图表到excel文件\n",
    "def chart_output(sorted_folder):\n",
    "    workbook = xlsxwriter.Workbook('./output/分析结果.xlsx')\n",
    "    worksheet = workbook.add_worksheet()\n",
    "    worksheet.write_row('A1', ['文件夹','占用空间(MB)'])\n",
    "    data = list(sorted_folder.items())\n",
    "    for row_num, (key, value) in enumerate(data, start=0):\n",
    "        worksheet.write(row_num+1, 0, key)   # 写入键（第一列）\n",
    "        worksheet.write(row_num+1, 1, value / (1024**2)) # 写入值（第二列）\n",
    "   \n",
    "    workbook.close()\n",
    "    \n",
    "def main(folder):\n",
    "    # 假设你有一个根目录或者一个包含多个子目录的列表\n",
    "    start_time = time.time()\n",
    "    result = scan_dirs(folder)\n",
    "    total_size = result[0]\n",
    "    statMap =  result[1]\n",
    "    end_time = time.time()\n",
    "    print(f\"{folder} Total size: {total_size / (1024 ** 3):.2f} GB, 分析耗时: {(end_time - start_time):.2f} 秒 \")\n",
    "    print(f\"开始打印瓶颈文件夹，总计 {len(statMap)}\")\n",
    "    # 打印瓶颈\n",
    "    sorted_folder = dict(sorted(statMap.items(), key=lambda item:item[1], reverse= True))    \n",
    "    for folder, size in sorted_folder.items():\n",
    "        print(f\"Folder: {folder}, size: {size / (1024**2):.2f} MB\")\n",
    "    chart_output(sorted_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c89bf4b2-b495-41de-980d-5f7e2e812efe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "开始分析 Folder: C: --------------------------------\n",
      "分析文件夹 C:$360Section 用时0.01 秒 ---------------------\n",
      "分析文件夹 C:$Recycle.Bin 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:$WinREAgent 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:AI_RecycleBin 用时0.16 秒 ---------------------\n",
      "分析文件夹 C:Documents and Settings 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:inetpub 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:Intel 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:KRECYCLE 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:MSOCache 用时0.04 秒 ---------------------\n",
      "分析文件夹 C:PerfLogs 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:Program Files 用时2.63 秒 ---------------------\n",
      "分析文件夹 C:Program Files (x86) 用时4.92 秒 ---------------------\n",
      "分析文件夹 C:ProgramData 用时1.31 秒 ---------------------\n",
      "分析文件夹 C:Recovery 用时0.01 秒 ---------------------\n",
      "分析文件夹 C:SnapPlugin 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:System Volume Information 用时0.00 秒 ---------------------\n",
      "分析文件夹 C:Users 用时22.08 秒 ---------------------\n",
      "分析文件夹 C:Windows 用时94.47 秒 ---------------------\n",
      "C: Total size: 55.68 GB, 分析耗时: 125.76 秒 \n",
      "开始打印瓶颈文件夹，总计 43\n",
      "Folder: C:Windows\\Installer, size: 2592.45 MB\n",
      "Folder: C:Windows\\System32, size: 2251.30 MB\n",
      "Folder: C:Windows\\SysWOW64, size: 1115.24 MB\n",
      "Folder: C:Windows\\System32\\DriverStore\\FileRepository\\nvcvegpu.inf_amd64_3bfacd0d07055664, size: 1045.16 MB\n",
      "Folder: C:Windows\\System32\\DriverStore\\FileRepository\\iigd_dch.inf_amd64_a1657f7761e46e74, size: 1025.46 MB\n",
      "Folder: C:Windows\\System32\\DriverStore\\FileRepository\\iigd_dch.inf_amd64_a086f01cc7be643a, size: 983.74 MB\n",
      "Folder: C:Program Files\\Docker\\Docker\\resources, size: 962.03 MB\n",
      "Folder: C:Windows\\Installer\\$PatchCache$\\Managed\\00005119110000000000000000F01FEC\\15.0.4569, size: 867.84 MB\n",
      "Folder: C:Windows\\Fonts, size: 590.40 MB\n",
      "Folder: C:Program Files\\Docker\\Docker\\resources\\wsl, size: 589.67 MB\n",
      "Folder: C:MSOCache\\All Users\\{91150000-0011-0000-0000-0000000FF1CE}-C, size: 551.61 MB\n",
      "Folder: C:Windows\\servicing\\LCU\\Package_for_RollupFix~31bf3856ad364e35~amd64~~19041.5965.1.5, size: 495.92 MB\n",
      "Folder: C:Program Files\\Git\\mingw64\\libexec\\git-core, size: 494.17 MB\n",
      "Folder: C:Program Files\\WSL, size: 434.88 MB\n",
      "Folder: C:Program Files (x86)\\Microsoft\\Edge\\Application\\138.0.3351.55, size: 434.26 MB\n",
      "Folder: C:Program Files (x86)\\Microsoft\\EdgeCore\\138.0.3351.55, size: 432.67 MB\n",
      "Folder: C:Program Files (x86)\\Microsoft\\EdgeCore\\137.0.3296.93, size: 430.74 MB\n",
      "Folder: C:Program Files (x86)\\Microsoft\\EdgeWebView\\Application\\137.0.3296.93, size: 430.74 MB\n",
      "Folder: C:Program Files (x86)\\Google\\GoogleUpdater\\crx_cache, size: 415.81 MB\n",
      "Folder: C:Program Files\\Docker\\cli-plugins, size: 411.82 MB\n",
      "Folder: C:Program Files\\Docker\\Docker\\resources\\cli-plugins, size: 411.82 MB\n",
      "Folder: C:Program Files (x86)\\Microsoft Office\\root\\Office16, size: 410.05 MB\n",
      "Folder: C:Program Files\\Google\\Chrome\\Application\\138.0.7204.96\\Installer, size: 409.09 MB\n",
      "Folder: C:Windows\\System32\\winevt\\Logs, size: 344.89 MB\n",
      "Folder: C:Program Files\\Google\\Chrome\\Application\\138.0.7204.96, size: 334.95 MB\n",
      "Folder: C:MSOCache\\All Users\\{90160000-003B-0000-0000-0000000FF1CE}-C, size: 289.75 MB\n",
      "Folder: C:Users\\luischen\\AppData\\Local\\Microsoft\\OneDrive\\24.244.1204.0003, size: 278.99 MB\n",
      "Folder: C:Program Files (x86)\\Microsoft Office\\Office15, size: 264.08 MB\n",
      "Folder: C:Users\\manulism\\AppData\\Local\\ConnectedDevicesPlatform\\8b8db738de9170f6, size: 243.48 MB\n",
      "Folder: C:Program Files\\PremiumSoft\\Navicat Premium 16, size: 242.43 MB\n",
      "Folder: C:Program Files\\Docker\\Docker\\resources\\bin, size: 237.55 MB\n",
      "Folder: C:Program Files (x86)\\Microsoft Office\\root\\vfs\\ProgramFilesCommonX86\\Microsoft Shared\\OFFICE16, size: 233.46 MB\n",
      "Folder: C:Program Files\\Docker\\Docker\\frontend, size: 227.81 MB\n",
      "Folder: C:Users\\manulism\\AppData\\Local\\Netease\\MailMaster\\web\\cache, size: 226.54 MB\n",
      "Folder: C:Program Files (x86)\\Microsoft Office\\root\\vfs\\Fonts\\private, size: 223.56 MB\n",
      "Folder: C:Program Files (x86)\\Windows Kits\\10\\Lib\\10.0.22621.0\\um\\arm64, size: 221.33 MB\n",
      "Folder: C:Users\\manulism\\AppData\\Local\\JetBrains\\IntelliJIdea2021.1\\index\\bytecodeanalysis, size: 219.57 MB\n",
      "Folder: C:Windows\\WinSxS\\amd64_windows-senseclient-service_31bf3856ad364e35_10.0.19041.5915_none_1c99fe0d44a3417e, size: 218.92 MB\n",
      "Folder: C:Users\\manulism\\AppData\\Local\\JetBrains\\IntelliJIdea2021.1\\index\\stubs, size: 214.40 MB\n",
      "Folder: C:Windows\\WinSxS\\Catalogs, size: 208.69 MB\n",
      "Folder: C:Program Files\\WSL\\tools, size: 207.56 MB\n",
      "Folder: C:Users\\manulism\\AppData\\Local\\JetBrains\\IntelliJIdea2021.1\\caches, size: 207.46 MB\n",
      "Folder: C:Users\\manulism\\AppData\\Roaming\\kdesktop\\Cache\\Cache_Data, size: 202.57 MB\n"
     ]
    }
   ],
   "source": [
    "folder = \"C:\"\n",
    "# %prun -s cumulative main(folder)\n",
    "main(folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c1c4469-fdf3-42bc-8a0e-da96624a273c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a3133c8-f757-4d94-9193-65ef0704010d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
